import sys
import os
import streamlit as st
from dotenv import load_dotenv

load_dotenv()
project_root = os.getenv("PROJECT_ROOT")
project_root_abs = os.path.abspath(project_root)
sys.path.append(project_root_abs)

from core.pipeline import NL2SQLPipeline


# 页面配置
st.set_page_config(page_title="NL2SQL", page_icon="")
st.markdown("""<style>.stApp {background-color: #f0f2f6;}<style>""", unsafe_allow_html=True)

# 初始化
if "pipeline" not in st.session_state:
    st.session_state.pipeline = NL2SQLPipeline()
if "history" not in st.session_state:
    st.session_state.history = []

# 页面标题
st.title("NL2SQL工具")
st.caption("输入数据查询需求，自动生成SQL并返回结果")

# 用户输入
user_query = st.text_input(
    "请输入查询请求",
    placeholder="在这里输入你的查询"
)

# 执行按钮
if st.button("生成并查询", type="primary") and user_query:
    with st.spinner("处理中..."):
        result_df, info = st.session_state.pipeline.run(user_query)

        # 显示SQL语句
        st.subheader("1. 生成的SQL")
        st.code(info["sql"], language="sql")

        # 显示校验结果
        st.subheader("2. SQL校验结果")
        if info["validation"]["is_valid"]:
            st.success(f"{info['validation']['valid_msg']}")
        else:
            st.error(f"{info['validation']['valid_msg']}")
            st.stop()

        # 显示查询结果
        st.subheader("3. 查询结果")
        if result_df is not None and not result_df.empty:
            st.dataframe(result_df, use_container_width=True)

            st.download_button(
                "下载结果（CSV）",
                data=result_df.to_csv(index=False),
                file_name="query_result.csv",
                mime="text/csv"
            )
        else:
            st.warning(f"查询无结果：{info['execution']}")

        # 保存历史
        st.session_state.history.append({
            "query": user_query,
            "sql": info["sql"],
            "success": result_df is not None
        })

# 显示历史
if st.session_state.history:
    with st.expander("查看历史", expanded=True):
        for i, item in enumerate(reversed(st.session_state.history), 1):
            st.write(f"{i}. {item['query']}")
            st.code(item["sql"], language="sql")
            status = "成功" if item["success"] else "失败"
            st.caption(f"状态：{status}")
            st.divider()
